import os
import zipfile
import tempfile
import shutil
from typing import List, Optional, Dict
import pymupdf
from docx import Document
from tqdm import tqdm
from utils.logger import logger
import pandas as pd
from config import Config

class DocumentLoader:
    def __init__(self):
        self.supported_extensions = {
            '.txt': self.load_single_txt,
            '.doc': self.load_single_docx,
            '.docx': self.load_single_docx,
            '.pdf': self.load_single_pdf,
            '.xlsx': self._load_excel,
            '.xls': self._load_excel,
            '.zip': self._load_zip
        }
        
    def _load_excel(self, file_path: str) -> Optional[str]:
        """加载Excel文件，处理标准问题、相似问题和标准答案
        
        预期的Excel格式:
        - 标准问题列: "标准问题"/"standard_question"
        - 相似问题列: "相似问题"/"similar_questions"
        - 标准答案列: "标准答案"/"answer"
        """
        try:
            df = pd.read_excel(file_path)
            
            # 统一列名
            column_mapping = {
                '标准问题': 'standard_question',
                'standard_question': 'standard_question',
                '相似问题': 'similar_questions',
                'similar_questions': 'similar_questions',
                '标准答案': 'answer',
                'answer': 'answer'
            }
            
            # 重命名列
            df = df.rename(columns=column_mapping)
            
            # 检查必要的列是否存在
            required_columns = ['standard_question', 'answer']
            if not all(col in df.columns for col in required_columns):
                logger.error(f"Excel文件缺少必要的列: {required_columns}")
                return None
            
            # 处理成QA对格式
            qa_pairs = []
            
            # 处理每一行
            for _, row in df.iterrows():
                # 处理标准问题
                standard_q = str(row['standard_question']).strip()
                if pd.isna(standard_q) or len(standard_q) < 2:
                    continue
                    
                answer = str(row['answer']).strip()
                if pd.isna(answer) or len(answer) < 2:
                    continue
                
                # 添加标准问题-答案对
                qa_pairs.append(f"{standard_q}\n答案：{answer}")
                
                # 处理相似问题（如果存在）
                if 'similar_questions' in df.columns:
                    similar_qs = str(row['similar_questions']).strip()
                    if not pd.isna(similar_qs) and len(similar_qs) > 3:
                        # 分割相似问题（支持多种分隔符）
                        similar_list = [q.strip() for q in similar_qs.replace('；', '\n').split('\n')]
                        # 添加每个相似问题-答案对
                        for sq in similar_list:
                            if len(sq) > 2:
                                qa_pairs.append(f"{sq}\n答案：{answer}")
            
            # 合并所有QA对
            return '####'.join(qa_pairs)
            
        except Exception as e:
            logger.error(f"加载Excel文件失败: {str(e)}")
            return None
    
    def _load_zip(self, file_path: str) -> Optional[str]:
        """加载ZIP文件，解压并处理所有支持的文档"""
        try:
            # 检查ZIP文件大小
            zip_size = os.path.getsize(file_path)
            if zip_size > Config.MAX_ZIP_SIZE:
                logger.error(f"ZIP文件过大: {zip_size} bytes > {Config.MAX_ZIP_SIZE} bytes")
                return None
            
            # 创建临时目录
            temp_dir = tempfile.mkdtemp()
            logger.info(f"创建临时目录: {temp_dir}")
            
            try:
                # 解压ZIP文件
                with zipfile.ZipFile(file_path, 'r') as zip_ref:
                    # 检查解压后的文件数量
                    file_list = zip_ref.namelist()
                    valid_files = [f for f in file_list if not f.startswith('__MACOSX/') and not f.endswith('/')]
                    
                    if len(valid_files) > Config.MAX_FILES_PER_BATCH:
                        logger.error(f"ZIP文件包含过多文件: {len(valid_files)} > {Config.MAX_FILES_PER_BATCH}")
                        return None
                    
                    # 解压所有文件
                    zip_ref.extractall(temp_dir)
                    logger.info(f"解压完成，包含 {len(valid_files)} 个文件")
                
                # 递归处理解压后的目录
                documents = self._load_directory_recursive(temp_dir)
                
                if not documents:
                    logger.warning("ZIP文件中没有找到有效文档")
                    return None
                
                # 合并所有文档内容
                all_content = []
                for doc_id, content in documents.items():
                    # 添加文档标识
                    all_content.append(f"### 文档: {doc_id} ###\n{content}")
                
                return "\n\n".join(all_content)
                
            finally:
                # 清理临时目录
                try:
                    shutil.rmtree(temp_dir)
                    logger.info(f"清理临时目录: {temp_dir}")
                except Exception as e:
                    logger.warning(f"清理临时目录失败: {str(e)}")
                    
        except Exception as e:
            logger.error(f"处理ZIP文件失败: {str(e)}")
            return None
    
    def _load_directory_recursive(self, directory_path: str) -> Dict[str, str]:
        """递归加载目录下的所有支持文档"""
        documents = {}
        
        try:
            for root, dirs, files in os.walk(directory_path):
                # 跳过系统目录
                dirs[:] = [d for d in dirs if not d.startswith('.') and d != '__MACOSX']
                
                for filename in files:
                    # 跳过隐藏文件和系统文件
                    if filename.startswith('.') or filename.startswith('~$'):
                        continue
                        
                    file_path = os.path.join(root, filename)
                    
                    # 检查文件扩展名
                    ext = os.path.splitext(filename)[1].lower()
                    if ext not in self.supported_extensions or ext == '.zip':  # 避免嵌套ZIP
                        continue
                    
                    try:
                        # 获取相对路径作为文档ID
                        rel_path = os.path.relpath(file_path, directory_path)
                        doc_id = rel_path.replace(os.sep, '_').replace(ext, '')
                        
                        # 加载文档内容
                        load_method = self.supported_extensions[ext]
                        content = load_method(file_path)
                        
                        if content and content.strip():
                            documents[doc_id] = content
                            logger.info(f"成功加载文档: {rel_path}")
                        else:
                            logger.warning(f"文档内容为空: {rel_path}")
                            
                    except Exception as e:
                        logger.error(f"加载文档失败 {filename}: {str(e)}")
                        continue
                        
        except Exception as e:
            logger.error(f"递归加载目录失败: {str(e)}")
            
        return documents
    
    def load_single_pdf(self, file_path: str) -> Optional[str]:
        """读取单个PDF文件并返回文本内容"""
        try:
            doc = pymupdf.open(file_path)
            text_content = []
            for page in doc:
                text = page.get_text("text")
                text_content.append(text)
            doc.close()
            return "\n".join(text_content)
        except Exception as e:
            logger.error(f"处理PDF文件 {file_path} 时发生错误: {str(e)}")
            return None
    
    def load_single_docx(self, file_path: str) -> Optional[str]:
        """读取单个Word文件并返回文本内容"""
        try:
            doc = Document(file_path)
            text_content = []
            for paragraph in doc.paragraphs:
                text_content.append(paragraph.text)
            return "\n".join(text_content)
        except Exception as e:
            logger.error(f"处理Word文件 {file_path} 时发生错误: {str(e)}")
            return None
    
    def load_single_txt(self, file_path: str) -> Optional[str]:
        """读取单个TXT文件并返回文本内容"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            try:
                # 尝试使用GBK编码
                with open(file_path, 'r', encoding='gbk') as f:
                    return f.read()
            except Exception as e:
                logger.error(f"处理TXT文件 {file_path} 时发生错误: {str(e)}")
                return None
        except Exception as e:
            logger.error(f"处理TXT文件 {file_path} 时发生错误: {str(e)}")
            return None
    
    def load_directory(self, directory_path: str) -> Dict[str, str]:
        """加载目录下的所有支持的文档"""
        try:
            logger.info(f"开始加载目录: {directory_path}")
            
            if not os.path.exists(directory_path):
                logger.error(f"目录不存在: {directory_path}")
                return {}
                
            documents = {}
            for filename in os.listdir(directory_path):
                file_path = os.path.join(directory_path, filename)
                logger.info(f"处理文件: {file_path}")
                
                if not os.path.isfile(file_path):
                    logger.warning(f"跳过非文件: {file_path}")
                    continue
                    
                ext = os.path.splitext(filename)[1].lower()
                if ext not in self.supported_extensions:
                    logger.warning(f"不支持的文件类型: {filename}")
                    continue
                    
                try:
                    # 使用映射的加载方法
                    load_method = self.supported_extensions[ext]
                    content = load_method(file_path)
                    
                    if content and content.strip():
                        # 使用文件名作为文档ID
                        doc_id = os.path.splitext(filename)[0]
                        documents[doc_id] = content
                        logger.info(f"成功加载文件: {filename}")
                    else:
                        logger.warning(f"文件内容为空: {filename}")
                        
                except Exception as e:
                    logger.error(f"处理文件失败 {filename}: {str(e)}", exc_info=True)
                    continue
                    
            logger.info(f"成功加载文档数量: {len(documents)}")
            return documents
            
        except Exception as e:
            logger.error(f"加载目录失败: {str(e)}", exc_info=True)
            return {}

    def is_valid_file(self, filename: str) -> bool:
        """检查文件是否有效"""
        if not filename:
            return False
        
        ext = os.path.splitext(filename)[1].lower()
        if ext not in self.supported_extensions:
            return False
        
        # 检查文件大小
        try:
            size = os.path.getsize(filename)
            if size > 10 * 1024 * 1024:  # 10MB
                logger.warning(f"文件过大: {filename}")
                return False
        except:
            return False
        
        return True